#!/usr/bin/python
# GM_MassDistrib.py
'''Create realistic populations of low-mass stars, up to a constant.'''

import numpy

import GM_InterpRoutines as IR

VariableNames = ['z', 'R', 'r', 'R0', 'hR', 'scale', 'MvMassScales', 'MvMass', 'massmults', 'massindices', 'delMvMass']

### This function is obsolete
def hz_from_Mv(simpars, phypars):
    for i in simpars:
        cmd = "%s = simpars['%s']" % (i,i)
        exec cmd
    for i in phypars:
        cmd = "%s = phypars['%s']" % (i,i)
        exec cmd
#    Mvdesired = MvMass[:,0]
#    Mv0 = numpy.array(Mvdesired, dtype=int)
#    Mv1 = Mv0 + 1
#    locs0 = numpy.equal(Mv0[:,numpy.newaxis], vanillaMvList[numpy.newaxis,:])#
#    locs1 = numpy.equal(Mv1[:,numpy.newaxis], vanillaMvList[numpy.newaxis,:])
#    Scale0 = numpy.multiply(locs0[:,:], vanillaScaleList[numpy.newaxis,:])
#    Scale1 = numpy.multiply(locs1[:,:], vanillaScaleList[numpy.newaxis,:])
#    Scale0 = Scale0.sum(axis=1)
#    Scale1 = Scale1.sum(axis=1)
#    (a,b) = IR.linear(Mv0,Scale0,Mv1,Scale1)
#    hz = a*Mvdesired + b
#    return hz


def MassDistrib(simpars, phypars):
    for i in simpars:
        cmd = "%s = simpars['%s']" % (i,i)
        exec cmd
    for i in phypars:
        cmd = "%s = phypars['%s']" % (i,i)
        exec cmd
    # R dependence of density, including a cutoff at r>scale to prevent assymetry about galactic center
    Rcorr = numpy.exp( -(R - R0)/hR )
    rmask = numpy.greater(scale, r)
    Rcorr = Rcorr*rmask
    distriblist = []
    dmlist = []
    rho0mlist = []
    del rmask
    # z dependence of density
    # The name change here is a relic of past design decisions and is not specially meaningful
#    hz = hz_from_Mv(simpars, phypars)
    hz = numpy.array(MvMassScales)
    for i in range(0, len(MvMass)):
        hzi = MvMassScales[i]
        m = MvMass[i,1]
        j0 = 9e9
        for j in range(0,3):
            if (m > masslimits[j,0]) and (m <= masslimits[j,1]):
                j0 = j
        alpha1 = massmults[j0]
        alpha2 = massindices[j0]
        dm = float(delMvMass[i,1])
#        if (i > 0) and (i < len(MvMass)-1):
#            mhi =  (MvMass[i+1,1] + MvMass[i,1])/2.
#            mlo =  (MvMass[i-1,1] + MvMass[i,1])/2.
#            dm = mhi - mlo
#        if i == 0:
#            mhi =  (MvMass[i+1,1] + MvMass[i,1])/2.
#            mlo =  MvMass[i,1]
#            dm = mhi - mlo
#        if i == len(MvMass)-1:
#            mhi =  MvMass[i,1]
#            mlo =  (MvMass[i-1,1] + MvMass[i,1])/2.
#            dm = mhi - mlo
        rho0m = alpha1*m**alpha2 # Massmults >0, massindindices <0; correct
        nm = rho0m*dm
        zcorr = numpy.exp(-z/hzi)
        thisdistrib = nm*numpy.multiply(Rcorr, zcorr[:,numpy.newaxis,:])
        distriblist.append(thisdistrib)
#        dmlist.append(dm)
        rho0mlist.append(rho0m)
    del Rcorr, zcorr
    return (distriblist, rho0mlist)
